"""
Regress one node on its parents with the Gaussian Mixture Model
"""
import numpy as np

from .base import GMMPredictor
from .prob_rf import ProbRF
from ..base import Regressor


class GMMRegressor(Regressor):
    """
    基于高斯混合模型的回归器
    
    使用高斯混合模型(GMM)来建模节点与其父节点之间的条件概率分布
    相比简单的线性回归，GMM可以捕捉更复杂的非线性关系和多模态分布
    """

    def __init__(
        self,
        regressor: GMMPredictor = None,  # GMM预测器实例
        sample_size: int = 100,          # 采样大小
        seed: int = None,                # 随机种子
        **kwargs
    ):
        """
        初始化GMM回归器
        
        Args:
            regressor: GMM预测器实例，默认使用概率随机森林(ProbRF)
            sample_size: 从GMM中采样的样本数量
            seed: 随机数生成器的种子
            **kwargs: 传递给父类的额外参数
        """
        super().__init__(**kwargs)
        # 如果没有指定预测器，则使用概率随机森林
        self._regressor = regressor if regressor else ProbRF(seed=seed)
        self._sample_size = sample_size

    def _score(
        self,
        train_x: np.ndarray,    # 训练集的父节点特征矩阵
        test_x: np.ndarray,     # 测试集的父节点特征矩阵
        train_y: np.ndarray,    # 训练集的当前节点值
        test_y: np.ndarray,     # 测试集的当前节点值
    ) -> np.ndarray:
        """
        计算异常分数
        
        工作流程：
        1. 使用训练数据训练GMM预测器
        2. 对每个测试样本：
           - 预测其条件概率分布（GMM）
           - 从分布中采样生成参考数据
           - 计算实际值相对于采样数据的Z分数
        
        Returns:
            np.ndarray: 每个测试样本的异常分数
        """
        # 训练GMM预测器
        self._regressor.train(train_x, train_y)
        # 对每个测试样本计算异常分数
        return np.array(
            [
                # 从预测的GMM中采样，并计算目标值的Z分数
                self._zscore(gmm.sample(self._sample_size), target)[0]
                for gmm, target in zip(self._regressor.predict(test_x), test_y)
            ]
        )
